import torch
import torch.nn as nn
import torch.quantization

# 定义一个简单的CNN模型
class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(20, 50, 5)
        self.fc1 = nn.Linear(50 * 4 * 4, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = self.pool(torch.relu(self.conv2(x)))
        x = x.view(-1, 50 * 4 * 4)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 实例化模型并加载预训练权重
model = SimpleCNN()
# 假设这里已经加载了预训练权重

# 准备量化配置
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')

# 融合模型中的模块以优化性能
torch.quantization.prepare(model, inplace=True)

# 量化模型
torch.quantization.convert(model, inplace=True)

# 现在model已经被量化，可以使用它进行推理
# ...（推理代码）
